# Copyright (C) 2024-2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import os
import sys

import tensorflow as tf

# Create the graph and model
class SampleGraph(tf.Module):
  def __init__(self):
    super(SampleGraph, self).__init__()
    self.var1 = tf.constant([[1, 0.75],[2000.43, -0.12345]], dtype=tf.float32)
  @tf.function(input_signature=[tf.TensorSpec([2,2], tf.float32)])
  def __call__(self, x):
    res = self.var1 + x
    return {'test_output_name': res}

module = SampleGraph()
sm_path = os.path.join(sys.argv[1], "dequantize")
tf.saved_model.save(module, sm_path)

converter = tf.lite.TFLiteConverter.from_saved_model(sm_path) # path to the SavedModel directory
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
]
converter.allow_custom_ops = True
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_model = converter.convert()

# Save the model.
with open(os.path.join(sys.argv[1], sm_path + ".tflite"), 'wb') as f:
  f.write(tflite_model)

# Model:
# (int8)tensor_0[12] -> DEQUANTIZE -> (float32)tensor_1[12]
# QUANTIZE:
#   scale: 0.25
#   zero point: 16
sm_path = os.path.join(sys.argv[1], "dequantize_int8")
tflite_model = b'\x08\x00\x00\x00TFL3\xbe\xff\xff\xff\x03\x00\x00\x00\x10\x00\x00\x00\x14\x00\x00\x00\x18\x00\x00\x00$\x00\x00\x00\x01\x00\x00\x00l\x01\x00\x00\x01\x00\x00\x00,\x00\x00\x00\n\x00\x00\x00test_model\x00\x00\x01\x00\x00\x00h\x00\x00\x00\x00\x00\x0e\x00\x18\x00\x04\x00\x08\x00\x0c\x00\x10\x00\x14\x00\x0e\x00\x00\x00\x14\x00\x00\x00\x1c\x00\x00\x00 \x00\x00\x00$\x00\x00\x00(\x00\x00\x00\x02\x00\x00\x00\xb4\x00\x00\x00\x80\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x000\x00\x00\x00\r\x00\x00\x00test_subgraph\x00\x00\x00\x04\x00\x06\x00\x04\x00\x00\x00\x00\x00\x0e\x00\x14\x00\x00\x00\x08\x00\x0c\x00\x07\x00\x10\x00\x0e\x00\x00\x00\x00\x00\x00&\x0c\x00\x00\x00\x10\x00\x00\x00\x18\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x04\x00\x04\x00\x04\x00\x00\x00\x0c\x00\x0c\x00\x04\x00\x00\x00\x00\x00\x08\x00\x0c\x00\x00\x00\x08\x00\x00\x00\x0c\x00\x00\x00\x01\x00\x00\x00\x0c\x00\x00\x00\x08\x00\x00\x00tensor_1\x00\x00\x0e\x00\x14\x00\x08\x00\x07\x00\x00\x00\x0c\x00\x10\x00\x0e\x00\x00\x00\x00\x00\x00\t\x0c\x00\x00\x00\x10\x00\x00\x00(\x00\x00\x00\x01\x00\x00\x00\x0c\x00\x00\x00\x08\x00\x00\x00tensor_0\x00\x00\x00\x00\x0c\x00\x0c\x00\x00\x00\x00\x00\x04\x00\x08\x00\x0c\x00\x00\x00\x08\x00\x00\x00\x0c\x00\x00\x00\x01\x00\x00\x00\x00\x00\x80>\x01\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x10\x00\x07\x00\x00\x00\x08\x00\x0c\x00\x0c\x00\x00\x00\x00\x00\x00\x06\x03\x00\x00\x00\x06\x00\x00\x00'
with open(os.path.join(sys.argv[1], sm_path + ".tflite"), 'wb') as f:
  f.write(tflite_model)

# Model:
# (uint8)tensor_0[12] -> DEQUANTIZE -> (float32)tensor_1[12]
# QUANTIZE:
#   scale: 0.25
#   zero point: 16
sm_path = os.path.join(sys.argv[1], "dequantize_uint8")
tflite_model = b'\x08\x00\x00\x00TFL3\xbe\xff\xff\xff\x03\x00\x00\x00\x10\x00\x00\x00\x14\x00\x00\x00\x18\x00\x00\x00$\x00\x00\x00\x01\x00\x00\x00l\x01\x00\x00\x01\x00\x00\x00,\x00\x00\x00\n\x00\x00\x00test_model\x00\x00\x01\x00\x00\x00h\x00\x00\x00\x00\x00\x0e\x00\x18\x00\x04\x00\x08\x00\x0c\x00\x10\x00\x14\x00\x0e\x00\x00\x00\x14\x00\x00\x00\x1c\x00\x00\x00 \x00\x00\x00$\x00\x00\x00(\x00\x00\x00\x02\x00\x00\x00\xb4\x00\x00\x00\x80\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x000\x00\x00\x00\r\x00\x00\x00test_subgraph\x00\x00\x00\x04\x00\x06\x00\x04\x00\x00\x00\x00\x00\x0e\x00\x14\x00\x00\x00\x08\x00\x0c\x00\x07\x00\x10\x00\x0e\x00\x00\x00\x00\x00\x00&\x0c\x00\x00\x00\x10\x00\x00\x00\x18\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x04\x00\x04\x00\x04\x00\x00\x00\x0c\x00\x0c\x00\x04\x00\x00\x00\x00\x00\x08\x00\x0c\x00\x00\x00\x08\x00\x00\x00\x0c\x00\x00\x00\x01\x00\x00\x00\x0c\x00\x00\x00\x08\x00\x00\x00tensor_1\x00\x00\x0e\x00\x14\x00\x08\x00\x07\x00\x00\x00\x0c\x00\x10\x00\x0e\x00\x00\x00\x00\x00\x00\x03\x0c\x00\x00\x00\x10\x00\x00\x00(\x00\x00\x00\x01\x00\x00\x00\x0c\x00\x00\x00\x08\x00\x00\x00tensor_0\x00\x00\x00\x00\x0c\x00\x0c\x00\x00\x00\x00\x00\x04\x00\x08\x00\x0c\x00\x00\x00\x08\x00\x00\x00\x0c\x00\x00\x00\x01\x00\x00\x00\x00\x00\x80>\x01\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0c\x00\x10\x00\x07\x00\x00\x00\x08\x00\x0c\x00\x0c\x00\x00\x00\x00\x00\x00\x06\x03\x00\x00\x00\x06\x00\x00\x00'
with open(os.path.join(sys.argv[1], sm_path + ".tflite"), 'wb') as f:
  f.write(tflite_model)
